/* 
 * bug h263 mounting util, (C) 2010 <neko@mikufans.cn> nekofs.
 *
 * This program is free software. It comes without any warranty, to
 * the extent permitted by applicable law. You can redistribute it
 * and/or modify it under the terms of the Do What The Fuck You Want
 * To Public License, Version 2, as published by Sam Hocevar. See
 * http://sam.zoy.org/wtfpl/COPYING for more details.
 *
 *
 * bug found by 真红の瞳 (http://blog.9acfun.cn/?author=5)
 * see http://blog.9acfun.cn/?p=194
 *
 * compile with:
 *   gcc mount.c -o mount.exe -Os -Wall -Wextra
 */ 

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "flv.h"
#if __MINGW32__
#  include <windows.h>
#  define sleep(n) Sleep(1000*n)
#endif

/**
 *  inject a crafted tag into flv stream
 *  @param f the output file stream
 *  @param ts timestamp of the injected tag
 *  @param keyframe whether this tag should be keyframe
 *  @param codec 0 for H263 by default, for other values refer to flv.h
 */
void inject_flvvideo_tag(FILE *f, int ts, int keyframe, int codec){
/* NB. This is a legal h263 packet. But now an empty tag works too, 
  so I comment this out.*/
/*  static char h263pkt[18] = 
      "\0\0\x84\0\x04\x04\x11\x27\xd7\xd7\xd7\xd7\xff\xde\0\0\0\x1a";
  static int h263pkt_s = 18; */
  int h263pkt_s = 0;
  u32 prevsize = bswap32(sizeof(flvtaghdr) + sizeof(flvvideohdr) + h263pkt_s);
  flvtaghdr fth = {
    .type = FLVTAG_VIDEO,
    .ds = bswap24(h263pkt_s + sizeof(flvvideohdr)),
    .ts = bswap24(ts & 0xffffff),
    .tse = ts >> 24,
    .id = 0,
  };
  flvvideohdr fvh = {
    .type = keyframe ? FLVVIDEO_KEY : FLVVIDEO_INTER,
    .codec = (codec ?: FLVVIDEO_H263),
  };
  fwrite(&fth, sizeof(fth), 1, f);
  fwrite(&fvh, sizeof(fvh), 1, f);
  //fwrite(h263pkt, 18, 1, f);
  fwrite(&prevsize, 4, 1, f);
}

void mount(FILE *fin, FILE *fout){
  char buf[BUFSIZ];
  flvhdr fh;
  flvtaghdr fth;
  size_t n, toread;
  int first;
  int ts;

  fseek(fin, 0, SEEK_SET);

  fread(&fh, sizeof(fh), 1, fin);
  fwrite(&fh, sizeof(fh), 1, fout);
  for(toread = bswap32(fh.doff) - sizeof(fh) + 4; toread > 0;){
    n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
    fwrite(buf, 1, n, fout);
    toread -= n;
  }

  first = 1;
  for(;;){
    if(!fread(&fth, sizeof(fth), 1, fin))
      break;

    if(fth.type != FLVTAG_VIDEO && fth.type != FLVTAG_AUDIO){
      fseek(fin, bswap24(fth.ds) + 4, SEEK_CUR);
      continue;
    }

    if(first){
      first = 0;
      inject_flvvideo_tag(fout, 0, 1, 0);
    }
    ts = bswap24(fth.ts) | (fth.tse << 24);
    ts += 360001;
    fth.ts = bswap24(ts & 0xffffff);
    fth.tse = ts >> 24;

    fwrite(&fth, sizeof(fth), 1, fout);
    for(toread = bswap24(fth.ds) + 4; toread > 0; toread -= n){
      n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
      fwrite(buf, n, 1, fout);
    }
  }
}
/**
 * detect and remove injected tags
 * @param fin input file stream
 * @param fout output file stream
 * @param vcodec vtag with codec != vcodec considered injected
 * @param acodec the same above
 *
 * plus, this function will detect timeline gap and remove it
 */
void unmount(FILE *fin, FILE *fout, int vcodec, int acodec){
  char buf[BUFSIZ];
  flvhdr fh;
  flvtaghdr fth;
  flvvideohdr fvh;
  flvaudiohdr fah;
  size_t n, toread;
  int ts;
  int framecount;
  int skip, last_ts, tsoff;

  fseek(fin, 0, SEEK_SET);
  fread(&fh, sizeof(fh), 1, fin);
  fwrite(&fh, sizeof(fh), 1, fout);
  for(toread = bswap32(fh.doff) - sizeof(fh) + 4; toread > 0;){
    n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
    fwrite(buf, 1, n, fout);
    toread -= n;
  }

  framecount = 0;
  skip = 0;
  last_ts = tsoff = 0;
  for(;;){
    if(!fread(&fth, sizeof(fth), 1, fin))
      break;

    ts = bswap24(fth.ts) | (fth.tse << 24);
    if(fth.type == FLVTAG_VIDEO){
      fread(&fvh, sizeof(fvh), 1, fin);
      fseek(fin, -sizeof(fvh), SEEK_CUR);
    }else if(fth.type == FLVTAG_AUDIO){
      fread(&fah, sizeof(fah), 1, fin);
      fseek(fin, -sizeof(fah), SEEK_CUR);
    }

    if((fth.type != FLVTAG_VIDEO && fth.type != FLVTAG_AUDIO)
    || (fth.type == FLVTAG_VIDEO && fvh.codec != vcodec)
    || (fth.type == FLVTAG_AUDIO && fah.format != acodec)
    || ts < last_ts){
      if(fth.type == FLVTAG_VIDEO || fth.type == FLVTAG_AUDIO)
        skip = 1;
      fprintf(stderr, "skip non-standard tag at %.3fs\n", ts/1e3);
      fseek(fin, bswap24(fth.ds) + 4, SEEK_CUR);
      continue;
    }

    //skip means there are non-standard tags before current position
    //gap means there is no non-standard tag before current position, 
    //but this tag and the last tag have noncontinuous timestamps
    if(skip){
      tsoff -= ts - last_ts - ( framecount ? (last_ts+tsoff)/framecount : 0 );
      fprintf(stderr, "non-standard tags before %.3fs, seek %.3fs\n", ts/1e3, tsoff/1e3);
      skip = 0;
    }else if(ts - last_ts > 1000){//XXX i set the gap threshold to 1000 msec
      tsoff -= ts - last_ts - ( framecount ? (last_ts+tsoff)/framecount : 0 );
      fprintf(stderr, "gap before %.3fs, seek %.3fs\n", ts/1e3, tsoff/1e3);
    }

    framecount++;
    last_ts = ts;

    ts += tsoff;
    fth.ts = bswap24(ts & 0xffffff);
    fth.tse = ts >> 24;


    fwrite(&fth, sizeof(fth), 1, fout);
    for(toread = bswap24(fth.ds)+4; toread > 0; toread -= n){
      n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
      fwrite(buf, n, 1, fout);
    }
  }
}

int is_mounted(FILE *fin, int *vcodec, int *acodec){
  flvtaghdr fth;
  flvvideohdr fvh;
  flvaudiohdr fah;
  int vc[16] = {0};
  int ac[16] = {0};
  int vcnum, acnum, i;
  int vmax, amax;

  fseek(fin, sizeof(flvhdr) + 4, SEEK_SET);

  for(;;){
    if(!fread(&fth, sizeof(fth), 1, fin))
      break;

    if(fth.type == FLVTAG_VIDEO){
      fread(&fvh, sizeof(fvh), 1, fin);
      fseek(fin, -sizeof(fvh), SEEK_CUR);
      vc[fvh.codec]++;
    }else if(fth.type == FLVTAG_AUDIO){
      fread(&fah, sizeof(fah), 1, fin);
      fseek(fin, -sizeof(fah), SEEK_CUR);
      ac[fah.format]++;
    }
    fseek(fin, bswap24(fth.ds) + 4, SEEK_CUR);
  }
  vcnum = acnum = vmax = amax = 0;
  for(i = 0; i < 16; i++){
    if(vc[i] > 0)vcnum++;
    if(ac[i] > 0)acnum++;
    if(vc[i] > vc[vmax])vmax = i;
    if(ac[i] > ac[amax])amax = i;
  }

  if(vcodec)*vcodec = vmax;
  if(acodec)*acodec = amax;
  return vcnum > 1 || acnum > 1;
}

int is_flv(FILE *fin){
  flvhdr fh;
  fseek(fin, 0, SEEK_SET);
  fread(&fh, sizeof(fh), 1, fin);
  return !strncmp((char*)fh.sign, "FLV", 3) && fh.has_video;
}

int main(int argc, char **argv){
  FILE *fin = NULL;
  FILE *fout = NULL;
  char *finame, *foname;
  int i, mounted;
  int vcnormal, acnormal;

  setbuf(stderr, NULL);
  for(i = 1; i < argc; i++){
    finame = argv[i];
    fprintf(stderr, "processing %s:\n", finame);
    if((fin = fopen(finame, "rb")) == NULL){
      fprintf(stderr, "can't open %s for reading\n", finame);
      sleep(5);
      return 1;
    }

    if(!is_flv(fin)){
      fprintf(stderr, "%s: not a valid flv file.\n", finame);
      sleep(1);
      fclose(fin);
      continue;
    }

    mounted = is_mounted(fin, &vcnormal, &acnormal);
    fprintf(stderr, "this file has %sbeen mounted.\n", mounted ? "" : "not ");

    foname = malloc(strlen(argv[1])+50+1);
    if(foname == NULL){
      fprintf(stderr, "malloc failed\n");
      sleep(5);
      fclose(fin);
      return 1;
    }
    strcpy(foname, finame);
    strcpy(foname + strlen(finame) - 4, mounted ? "-unmount.flv" : "-mount.flv");

    fout = fopen(foname, "wb");
    free(foname);
    if(fout == NULL){
      fprintf(stderr, "can't open %s for writing\n", foname);
      sleep(5);
      fclose(fin);
      return 1;
    }

    if(mounted)
      unmount(fin, fout, vcnormal, acnormal);
    else
      mount(fin, fout);

    fclose(fin);
    fclose(fout);
  }
  fprintf(stderr, "finished.\n");
  sleep(2);
  return 0;
}
